import torch
from transformers import BertTokenizer, BertForSequenceClassification

def test_model(texts):
    """
    使用训练好的模型对输入文本进行分类预测
    :param texts: 待预测的文本列表
    """
    # 加载保存的分词器
    tokenizer = BertTokenizer.from_pretrained('weights/sort_model')
    # 加载保存的模型
    model = BertForSequenceClassification.from_pretrained('weights/sort_model')

    # 将模型移动到合适的设备上（如果有 GPU 则使用 GPU，否则使用 CPU）
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    model.eval()

    label_map = {0: '技术问题', 1: '售前咨询', 2: '投诉', 3: '聊天'}

    with torch.no_grad():
        for text in texts:
            # 对文本进行分词和编码
            encoding = tokenizer.encode_plus(
                text,
                add_special_tokens=True,
                max_length=128,
                truncation=True,
                padding='max_length',
                return_attention_mask=True,
                return_tensors='pt'
            )
            input_ids = encoding['input_ids'].to(device)
            attention_mask = encoding['attention_mask'].to(device)

            # 使用模型进行预测
            outputs = model(input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            _, prediction = torch.max(logits, dim=1)

            predicted_label = label_map[prediction.item()]
            print(f"文本: {text} 被预测为: {predicted_label}")

if __name__ == "__main__":
    sample_texts = [
        "我在连接网络的时候遇到了问题，一直显示连接失败",
        "想咨询一下你们产品的价格和功能",
        "客服的服务态度特别差，我很不满意",
        "今天天气真好，随便聊聊"
    ]
    test_model(sample_texts)